# State dict keys and shapes for an InstantX FLUX ControlNet Union model. Intended to be used for unit tests.
# These keys were extracted from:
# https://huggingface.co/InstantX/FLUX.1-dev-Controlnet-Union/blob/4f32d6f2b220f8873d49bb8acc073e1df180c994/diffusion_pytorch_model.safetensors
instantx_sd_shapes = {
    "context_embedder.bias": [3072],
    "context_embedder.weight": [3072, 4096],
    "controlnet_blocks.0.bias": [3072],
    "controlnet_blocks.0.weight": [3072, 3072],
    "controlnet_blocks.1.bias": [3072],
    "controlnet_blocks.1.weight": [3072, 3072],
    "controlnet_blocks.2.bias": [3072],
    "controlnet_blocks.2.weight": [3072, 3072],
    "controlnet_blocks.3.bias": [3072],
    "controlnet_blocks.3.weight": [3072, 3072],
    "controlnet_blocks.4.bias": [3072],
    "controlnet_blocks.4.weight": [3072, 3072],
    "controlnet_mode_embedder.weight": [10, 3072],
    "controlnet_single_blocks.0.bias": [3072],
    "controlnet_single_blocks.0.weight": [3072, 3072],
    "controlnet_single_blocks.1.bias": [3072],
    "controlnet_single_blocks.1.weight": [3072, 3072],
    "controlnet_single_blocks.2.bias": [3072],
    "controlnet_single_blocks.2.weight": [3072, 3072],
    "controlnet_single_blocks.3.bias": [3072],
    "controlnet_single_blocks.3.weight": [3072, 3072],
    "controlnet_single_blocks.4.bias": [3072],
    "controlnet_single_blocks.4.weight": [3072, 3072],
    "controlnet_single_blocks.5.bias": [3072],
    "controlnet_single_blocks.5.weight": [3072, 3072],
    "controlnet_single_blocks.6.bias": [3072],
    "controlnet_single_blocks.6.weight": [3072, 3072],
    "controlnet_single_blocks.7.bias": [3072],
    "controlnet_single_blocks.7.weight": [3072, 3072],
    "controlnet_single_blocks.8.bias": [3072],
    "controlnet_single_blocks.8.weight": [3072, 3072],
    "controlnet_single_blocks.9.bias": [3072],
    "controlnet_single_blocks.9.weight": [3072, 3072],
    "controlnet_x_embedder.bias": [3072],
    "controlnet_x_embedder.weight": [3072, 64],
    "single_transformer_blocks.0.attn.norm_k.weight": [128],
    "single_transformer_blocks.0.attn.norm_q.weight": [128],
    "single_transformer_blocks.0.attn.to_k.bias": [3072],
    "single_transformer_blocks.0.attn.to_k.weight": [3072, 3072],
    "single_transformer_blocks.0.attn.to_q.bias": [3072],
    "single_transformer_blocks.0.attn.to_q.weight": [3072, 3072],
    "single_transformer_blocks.0.attn.to_v.bias": [3072],
    "single_transformer_blocks.0.attn.to_v.weight": [3072, 3072],
    "single_transformer_blocks.0.norm.linear.bias": [9216],
    "single_transformer_blocks.0.norm.linear.weight": [9216, 3072],
    "single_transformer_blocks.0.proj_mlp.bias": [12288],
    "single_transformer_blocks.0.proj_mlp.weight": [12288, 3072],
    "single_transformer_blocks.0.proj_out.bias": [3072],
    "single_transformer_blocks.0.proj_out.weight": [3072, 15360],
    "single_transformer_blocks.1.attn.norm_k.weight": [128],
    "single_transformer_blocks.1.attn.norm_q.weight": [128],
    "single_transformer_blocks.1.attn.to_k.bias": [3072],
    "single_transformer_blocks.1.attn.to_k.weight": [3072, 3072],
    "single_transformer_blocks.1.attn.to_q.bias": [3072],
    "single_transformer_blocks.1.attn.to_q.weight": [3072, 3072],
    "single_transformer_blocks.1.attn.to_v.bias": [3072],
    "single_transformer_blocks.1.attn.to_v.weight": [3072, 3072],
    "single_transformer_blocks.1.norm.linear.bias": [9216],
    "single_transformer_blocks.1.norm.linear.weight": [9216, 3072],
    "single_transformer_blocks.1.proj_mlp.bias": [12288],
    "single_transformer_blocks.1.proj_mlp.weight": [12288, 3072],
    "single_transformer_blocks.1.proj_out.bias": [3072],
    "single_transformer_blocks.1.proj_out.weight": [3072, 15360],
    "single_transformer_blocks.2.attn.norm_k.weight": [128],
    "single_transformer_blocks.2.attn.norm_q.weight": [128],
    "single_transformer_blocks.2.attn.to_k.bias": [3072],
    "single_transformer_blocks.2.attn.to_k.weight": [3072, 3072],
    "single_transformer_blocks.2.attn.to_q.bias": [3072],
    "single_transformer_blocks.2.attn.to_q.weight": [3072, 3072],
    "single_transformer_blocks.2.attn.to_v.bias": [3072],
    "single_transformer_blocks.2.attn.to_v.weight": [3072, 3072],
    "single_transformer_blocks.2.norm.linear.bias": [9216],
    "single_transformer_blocks.2.norm.linear.weight": [9216, 3072],
    "single_transformer_blocks.2.proj_mlp.bias": [12288],
    "single_transformer_blocks.2.proj_mlp.weight": [12288, 3072],
    "single_transformer_blocks.2.proj_out.bias": [3072],
    "single_transformer_blocks.2.proj_out.weight": [3072, 15360],
    "single_transformer_blocks.3.attn.norm_k.weight": [128],
    "single_transformer_blocks.3.attn.norm_q.weight": [128],
    "single_transformer_blocks.3.attn.to_k.bias": [3072],
    "single_transformer_blocks.3.attn.to_k.weight": [3072, 3072],
    "single_transformer_blocks.3.attn.to_q.bias": [3072],
    "single_transformer_blocks.3.attn.to_q.weight": [3072, 3072],
    "single_transformer_blocks.3.attn.to_v.bias": [3072],
    "single_transformer_blocks.3.attn.to_v.weight": [3072, 3072],
    "single_transformer_blocks.3.norm.linear.bias": [9216],
    "single_transformer_blocks.3.norm.linear.weight": [9216, 3072],
    "single_transformer_blocks.3.proj_mlp.bias": [12288],
    "single_transformer_blocks.3.proj_mlp.weight": [12288, 3072],
    "single_transformer_blocks.3.proj_out.bias": [3072],
    "single_transformer_blocks.3.proj_out.weight": [3072, 15360],
    "single_transformer_blocks.4.attn.norm_k.weight": [128],
    "single_transformer_blocks.4.attn.norm_q.weight": [128],
    "single_transformer_blocks.4.attn.to_k.bias": [3072],
    "single_transformer_blocks.4.attn.to_k.weight": [3072, 3072],
    "single_transformer_blocks.4.attn.to_q.bias": [3072],
    "single_transformer_blocks.4.attn.to_q.weight": [3072, 3072],
    "single_transformer_blocks.4.attn.to_v.bias": [3072],
    "single_transformer_blocks.4.attn.to_v.weight": [3072, 3072],
    "single_transformer_blocks.4.norm.linear.bias": [9216],
    "single_transformer_blocks.4.norm.linear.weight": [9216, 3072],
    "single_transformer_blocks.4.proj_mlp.bias": [12288],
    "single_transformer_blocks.4.proj_mlp.weight": [12288, 3072],
    "single_transformer_blocks.4.proj_out.bias": [3072],
    "single_transformer_blocks.4.proj_out.weight": [3072, 15360],
    "single_transformer_blocks.5.attn.norm_k.weight": [128],
    "single_transformer_blocks.5.attn.norm_q.weight": [128],
    "single_transformer_blocks.5.attn.to_k.bias": [3072],
    "single_transformer_blocks.5.attn.to_k.weight": [3072, 3072],
    "single_transformer_blocks.5.attn.to_q.bias": [3072],
    "single_transformer_blocks.5.attn.to_q.weight": [3072, 3072],
    "single_transformer_blocks.5.attn.to_v.bias": [3072],
    "single_transformer_blocks.5.attn.to_v.weight": [3072, 3072],
    "single_transformer_blocks.5.norm.linear.bias": [9216],
    "single_transformer_blocks.5.norm.linear.weight": [9216, 3072],
    "single_transformer_blocks.5.proj_mlp.bias": [12288],
    "single_transformer_blocks.5.proj_mlp.weight": [12288, 3072],
    "single_transformer_blocks.5.proj_out.bias": [3072],
    "single_transformer_blocks.5.proj_out.weight": [3072, 15360],
    "single_transformer_blocks.6.attn.norm_k.weight": [128],
    "single_transformer_blocks.6.attn.norm_q.weight": [128],
    "single_transformer_blocks.6.attn.to_k.bias": [3072],
    "single_transformer_blocks.6.attn.to_k.weight": [3072, 3072],
    "single_transformer_blocks.6.attn.to_q.bias": [3072],
    "single_transformer_blocks.6.attn.to_q.weight": [3072, 3072],
    "single_transformer_blocks.6.attn.to_v.bias": [3072],
    "single_transformer_blocks.6.attn.to_v.weight": [3072, 3072],
    "single_transformer_blocks.6.norm.linear.bias": [9216],
    "single_transformer_blocks.6.norm.linear.weight": [9216, 3072],
    "single_transformer_blocks.6.proj_mlp.bias": [12288],
    "single_transformer_blocks.6.proj_mlp.weight": [12288, 3072],
    "single_transformer_blocks.6.proj_out.bias": [3072],
    "single_transformer_blocks.6.proj_out.weight": [3072, 15360],
    "single_transformer_blocks.7.attn.norm_k.weight": [128],
    "single_transformer_blocks.7.attn.norm_q.weight": [128],
    "single_transformer_blocks.7.attn.to_k.bias": [3072],
    "single_transformer_blocks.7.attn.to_k.weight": [3072, 3072],
    "single_transformer_blocks.7.attn.to_q.bias": [3072],
    "single_transformer_blocks.7.attn.to_q.weight": [3072, 3072],
    "single_transformer_blocks.7.attn.to_v.bias": [3072],
    "single_transformer_blocks.7.attn.to_v.weight": [3072, 3072],
    "single_transformer_blocks.7.norm.linear.bias": [9216],
    "single_transformer_blocks.7.norm.linear.weight": [9216, 3072],
    "single_transformer_blocks.7.proj_mlp.bias": [12288],
    "single_transformer_blocks.7.proj_mlp.weight": [12288, 3072],
    "single_transformer_blocks.7.proj_out.bias": [3072],
    "single_transformer_blocks.7.proj_out.weight": [3072, 15360],
    "single_transformer_blocks.8.attn.norm_k.weight": [128],
    "single_transformer_blocks.8.attn.norm_q.weight": [128],
    "single_transformer_blocks.8.attn.to_k.bias": [3072],
    "single_transformer_blocks.8.attn.to_k.weight": [3072, 3072],
    "single_transformer_blocks.8.attn.to_q.bias": [3072],
    "single_transformer_blocks.8.attn.to_q.weight": [3072, 3072],
    "single_transformer_blocks.8.attn.to_v.bias": [3072],
    "single_transformer_blocks.8.attn.to_v.weight": [3072, 3072],
    "single_transformer_blocks.8.norm.linear.bias": [9216],
    "single_transformer_blocks.8.norm.linear.weight": [9216, 3072],
    "single_transformer_blocks.8.proj_mlp.bias": [12288],
    "single_transformer_blocks.8.proj_mlp.weight": [12288, 3072],
    "single_transformer_blocks.8.proj_out.bias": [3072],
    "single_transformer_blocks.8.proj_out.weight": [3072, 15360],
    "single_transformer_blocks.9.attn.norm_k.weight": [128],
    "single_transformer_blocks.9.attn.norm_q.weight": [128],
    "single_transformer_blocks.9.attn.to_k.bias": [3072],
    "single_transformer_blocks.9.attn.to_k.weight": [3072, 3072],
    "single_transformer_blocks.9.attn.to_q.bias": [3072],
    "single_transformer_blocks.9.attn.to_q.weight": [3072, 3072],
    "single_transformer_blocks.9.attn.to_v.bias": [3072],
    "single_transformer_blocks.9.attn.to_v.weight": [3072, 3072],
    "single_transformer_blocks.9.norm.linear.bias": [9216],
    "single_transformer_blocks.9.norm.linear.weight": [9216, 3072],
    "single_transformer_blocks.9.proj_mlp.bias": [12288],
    "single_transformer_blocks.9.proj_mlp.weight": [12288, 3072],
    "single_transformer_blocks.9.proj_out.bias": [3072],
    "single_transformer_blocks.9.proj_out.weight": [3072, 15360],
    "time_text_embed.guidance_embedder.linear_1.bias": [3072],
    "time_text_embed.guidance_embedder.linear_1.weight": [3072, 256],
    "time_text_embed.guidance_embedder.linear_2.bias": [3072],
    "time_text_embed.guidance_embedder.linear_2.weight": [3072, 3072],
    "time_text_embed.text_embedder.linear_1.bias": [3072],
    "time_text_embed.text_embedder.linear_1.weight": [3072, 768],
    "time_text_embed.text_embedder.linear_2.bias": [3072],
    "time_text_embed.text_embedder.linear_2.weight": [3072, 3072],
    "time_text_embed.timestep_embedder.linear_1.bias": [3072],
    "time_text_embed.timestep_embedder.linear_1.weight": [3072, 256],
    "time_text_embed.timestep_embedder.linear_2.bias": [3072],
    "time_text_embed.timestep_embedder.linear_2.weight": [3072, 3072],
    "transformer_blocks.0.attn.add_k_proj.bias": [3072],
    "transformer_blocks.0.attn.add_k_proj.weight": [3072, 3072],
    "transformer_blocks.0.attn.add_q_proj.bias": [3072],
    "transformer_blocks.0.attn.add_q_proj.weight": [3072, 3072],
    "transformer_blocks.0.attn.add_v_proj.bias": [3072],
    "transformer_blocks.0.attn.add_v_proj.weight": [3072, 3072],
    "transformer_blocks.0.attn.norm_added_k.weight": [128],
    "transformer_blocks.0.attn.norm_added_q.weight": [128],
    "transformer_blocks.0.attn.norm_k.weight": [128],
    "transformer_blocks.0.attn.norm_q.weight": [128],
    "transformer_blocks.0.attn.to_add_out.bias": [3072],
    "transformer_blocks.0.attn.to_add_out.weight": [3072, 3072],
    "transformer_blocks.0.attn.to_k.bias": [3072],
    "transformer_blocks.0.attn.to_k.weight": [3072, 3072],
    "transformer_blocks.0.attn.to_out.0.bias": [3072],
    "transformer_blocks.0.attn.to_out.0.weight": [3072, 3072],
    "transformer_blocks.0.attn.to_q.bias": [3072],
    "transformer_blocks.0.attn.to_q.weight": [3072, 3072],
    "transformer_blocks.0.attn.to_v.bias": [3072],
    "transformer_blocks.0.attn.to_v.weight": [3072, 3072],
    "transformer_blocks.0.ff.net.0.proj.bias": [12288],
    "transformer_blocks.0.ff.net.0.proj.weight": [12288, 3072],
    "transformer_blocks.0.ff.net.2.bias": [3072],
    "transformer_blocks.0.ff.net.2.weight": [3072, 12288],
    "transformer_blocks.0.ff_context.net.0.proj.bias": [12288],
    "transformer_blocks.0.ff_context.net.0.proj.weight": [12288, 3072],
    "transformer_blocks.0.ff_context.net.2.bias": [3072],
    "transformer_blocks.0.ff_context.net.2.weight": [3072, 12288],
    "transformer_blocks.0.norm1.linear.bias": [18432],
    "transformer_blocks.0.norm1.linear.weight": [18432, 3072],
    "transformer_blocks.0.norm1_context.linear.bias": [18432],
    "transformer_blocks.0.norm1_context.linear.weight": [18432, 3072],
    "transformer_blocks.1.attn.add_k_proj.bias": [3072],
    "transformer_blocks.1.attn.add_k_proj.weight": [3072, 3072],
    "transformer_blocks.1.attn.add_q_proj.bias": [3072],
    "transformer_blocks.1.attn.add_q_proj.weight": [3072, 3072],
    "transformer_blocks.1.attn.add_v_proj.bias": [3072],
    "transformer_blocks.1.attn.add_v_proj.weight": [3072, 3072],
    "transformer_blocks.1.attn.norm_added_k.weight": [128],
    "transformer_blocks.1.attn.norm_added_q.weight": [128],
    "transformer_blocks.1.attn.norm_k.weight": [128],
    "transformer_blocks.1.attn.norm_q.weight": [128],
    "transformer_blocks.1.attn.to_add_out.bias": [3072],
    "transformer_blocks.1.attn.to_add_out.weight": [3072, 3072],
    "transformer_blocks.1.attn.to_k.bias": [3072],
    "transformer_blocks.1.attn.to_k.weight": [3072, 3072],
    "transformer_blocks.1.attn.to_out.0.bias": [3072],
    "transformer_blocks.1.attn.to_out.0.weight": [3072, 3072],
    "transformer_blocks.1.attn.to_q.bias": [3072],
    "transformer_blocks.1.attn.to_q.weight": [3072, 3072],
    "transformer_blocks.1.attn.to_v.bias": [3072],
    "transformer_blocks.1.attn.to_v.weight": [3072, 3072],
    "transformer_blocks.1.ff.net.0.proj.bias": [12288],
    "transformer_blocks.1.ff.net.0.proj.weight": [12288, 3072],
    "transformer_blocks.1.ff.net.2.bias": [3072],
    "transformer_blocks.1.ff.net.2.weight": [3072, 12288],
    "transformer_blocks.1.ff_context.net.0.proj.bias": [12288],
    "transformer_blocks.1.ff_context.net.0.proj.weight": [12288, 3072],
    "transformer_blocks.1.ff_context.net.2.bias": [3072],
    "transformer_blocks.1.ff_context.net.2.weight": [3072, 12288],
    "transformer_blocks.1.norm1.linear.bias": [18432],
    "transformer_blocks.1.norm1.linear.weight": [18432, 3072],
    "transformer_blocks.1.norm1_context.linear.bias": [18432],
    "transformer_blocks.1.norm1_context.linear.weight": [18432, 3072],
    "transformer_blocks.2.attn.add_k_proj.bias": [3072],
    "transformer_blocks.2.attn.add_k_proj.weight": [3072, 3072],
    "transformer_blocks.2.attn.add_q_proj.bias": [3072],
    "transformer_blocks.2.attn.add_q_proj.weight": [3072, 3072],
    "transformer_blocks.2.attn.add_v_proj.bias": [3072],
    "transformer_blocks.2.attn.add_v_proj.weight": [3072, 3072],
    "transformer_blocks.2.attn.norm_added_k.weight": [128],
    "transformer_blocks.2.attn.norm_added_q.weight": [128],
    "transformer_blocks.2.attn.norm_k.weight": [128],
    "transformer_blocks.2.attn.norm_q.weight": [128],
    "transformer_blocks.2.attn.to_add_out.bias": [3072],
    "transformer_blocks.2.attn.to_add_out.weight": [3072, 3072],
    "transformer_blocks.2.attn.to_k.bias": [3072],
    "transformer_blocks.2.attn.to_k.weight": [3072, 3072],
    "transformer_blocks.2.attn.to_out.0.bias": [3072],
    "transformer_blocks.2.attn.to_out.0.weight": [3072, 3072],
    "transformer_blocks.2.attn.to_q.bias": [3072],
    "transformer_blocks.2.attn.to_q.weight": [3072, 3072],
    "transformer_blocks.2.attn.to_v.bias": [3072],
    "transformer_blocks.2.attn.to_v.weight": [3072, 3072],
    "transformer_blocks.2.ff.net.0.proj.bias": [12288],
    "transformer_blocks.2.ff.net.0.proj.weight": [12288, 3072],
    "transformer_blocks.2.ff.net.2.bias": [3072],
    "transformer_blocks.2.ff.net.2.weight": [3072, 12288],
    "transformer_blocks.2.ff_context.net.0.proj.bias": [12288],
    "transformer_blocks.2.ff_context.net.0.proj.weight": [12288, 3072],
    "transformer_blocks.2.ff_context.net.2.bias": [3072],
    "transformer_blocks.2.ff_context.net.2.weight": [3072, 12288],
    "transformer_blocks.2.norm1.linear.bias": [18432],
    "transformer_blocks.2.norm1.linear.weight": [18432, 3072],
    "transformer_blocks.2.norm1_context.linear.bias": [18432],
    "transformer_blocks.2.norm1_context.linear.weight": [18432, 3072],
    "transformer_blocks.3.attn.add_k_proj.bias": [3072],
    "transformer_blocks.3.attn.add_k_proj.weight": [3072, 3072],
    "transformer_blocks.3.attn.add_q_proj.bias": [3072],
    "transformer_blocks.3.attn.add_q_proj.weight": [3072, 3072],
    "transformer_blocks.3.attn.add_v_proj.bias": [3072],
    "transformer_blocks.3.attn.add_v_proj.weight": [3072, 3072],
    "transformer_blocks.3.attn.norm_added_k.weight": [128],
    "transformer_blocks.3.attn.norm_added_q.weight": [128],
    "transformer_blocks.3.attn.norm_k.weight": [128],
    "transformer_blocks.3.attn.norm_q.weight": [128],
    "transformer_blocks.3.attn.to_add_out.bias": [3072],
    "transformer_blocks.3.attn.to_add_out.weight": [3072, 3072],
    "transformer_blocks.3.attn.to_k.bias": [3072],
    "transformer_blocks.3.attn.to_k.weight": [3072, 3072],
    "transformer_blocks.3.attn.to_out.0.bias": [3072],
    "transformer_blocks.3.attn.to_out.0.weight": [3072, 3072],
    "transformer_blocks.3.attn.to_q.bias": [3072],
    "transformer_blocks.3.attn.to_q.weight": [3072, 3072],
    "transformer_blocks.3.attn.to_v.bias": [3072],
    "transformer_blocks.3.attn.to_v.weight": [3072, 3072],
    "transformer_blocks.3.ff.net.0.proj.bias": [12288],
    "transformer_blocks.3.ff.net.0.proj.weight": [12288, 3072],
    "transformer_blocks.3.ff.net.2.bias": [3072],
    "transformer_blocks.3.ff.net.2.weight": [3072, 12288],
    "transformer_blocks.3.ff_context.net.0.proj.bias": [12288],
    "transformer_blocks.3.ff_context.net.0.proj.weight": [12288, 3072],
    "transformer_blocks.3.ff_context.net.2.bias": [3072],
    "transformer_blocks.3.ff_context.net.2.weight": [3072, 12288],
    "transformer_blocks.3.norm1.linear.bias": [18432],
    "transformer_blocks.3.norm1.linear.weight": [18432, 3072],
    "transformer_blocks.3.norm1_context.linear.bias": [18432],
    "transformer_blocks.3.norm1_context.linear.weight": [18432, 3072],
    "transformer_blocks.4.attn.add_k_proj.bias": [3072],
    "transformer_blocks.4.attn.add_k_proj.weight": [3072, 3072],
    "transformer_blocks.4.attn.add_q_proj.bias": [3072],
    "transformer_blocks.4.attn.add_q_proj.weight": [3072, 3072],
    "transformer_blocks.4.attn.add_v_proj.bias": [3072],
    "transformer_blocks.4.attn.add_v_proj.weight": [3072, 3072],
    "transformer_blocks.4.attn.norm_added_k.weight": [128],
    "transformer_blocks.4.attn.norm_added_q.weight": [128],
    "transformer_blocks.4.attn.norm_k.weight": [128],
    "transformer_blocks.4.attn.norm_q.weight": [128],
    "transformer_blocks.4.attn.to_add_out.bias": [3072],
    "transformer_blocks.4.attn.to_add_out.weight": [3072, 3072],
    "transformer_blocks.4.attn.to_k.bias": [3072],
    "transformer_blocks.4.attn.to_k.weight": [3072, 3072],
    "transformer_blocks.4.attn.to_out.0.bias": [3072],
    "transformer_blocks.4.attn.to_out.0.weight": [3072, 3072],
    "transformer_blocks.4.attn.to_q.bias": [3072],
    "transformer_blocks.4.attn.to_q.weight": [3072, 3072],
    "transformer_blocks.4.attn.to_v.bias": [3072],
    "transformer_blocks.4.attn.to_v.weight": [3072, 3072],
    "transformer_blocks.4.ff.net.0.proj.bias": [12288],
    "transformer_blocks.4.ff.net.0.proj.weight": [12288, 3072],
    "transformer_blocks.4.ff.net.2.bias": [3072],
    "transformer_blocks.4.ff.net.2.weight": [3072, 12288],
    "transformer_blocks.4.ff_context.net.0.proj.bias": [12288],
    "transformer_blocks.4.ff_context.net.0.proj.weight": [12288, 3072],
    "transformer_blocks.4.ff_context.net.2.bias": [3072],
    "transformer_blocks.4.ff_context.net.2.weight": [3072, 12288],
    "transformer_blocks.4.norm1.linear.bias": [18432],
    "transformer_blocks.4.norm1.linear.weight": [18432, 3072],
    "transformer_blocks.4.norm1_context.linear.bias": [18432],
    "transformer_blocks.4.norm1_context.linear.weight": [18432, 3072],
    "x_embedder.bias": [3072],
    "x_embedder.weight": [3072, 64],
}


# InstantX FLUX ControlNet config for unit tests.
# Copied from https://huggingface.co/InstantX/FLUX.1-dev-Controlnet-Union/blob/main/config.json
instantx_config = {
    "_class_name": "FluxControlNetModel",
    "_diffusers_version": "0.30.0.dev0",
    "_name_or_path": "/mnt/wangqixun/",
    "attention_head_dim": 128,
    "axes_dims_rope": [16, 56, 56],
    "guidance_embeds": True,
    "in_channels": 64,
    "joint_attention_dim": 4096,
    "num_attention_heads": 24,
    "num_layers": 5,
    "num_mode": 10,
    "num_single_layers": 10,
    "patch_size": 1,
    "pooled_projection_dim": 768,
}
